import os
import sys
from PIL import Image
import random
from torch.utils.data.dataset import Dataset
from transforms.x_transform import *
import torchvision

MNIST_TRAIN_PATH = r"E:\deep_learning_tutorial\datasets\MNIST\train"
MNIST_TEST_PATH = r"E:\deep_learning_tutorial\datasets\MNIST\test"

'''
MNIST手写数字数据集
训练集共60000张图像
测试集共10000张图像
'''

class MNISTRecord(object):
    '''
    MNIST数据集的每个样本的记录，存储样本的相关信息
    '''
    def __init__(self, path):
        '''
        根据给定的路径字符串，对样本的属性进行解析
        :param path: 样本路径
        '''
        assert os.path.exists(path), path
        self.path = path
        # 图像所在的文件夹名称即为图像的标签
        self.category = int(os.path.split(self.path)[0][-1])


class MNIST(Dataset):
    def __init__(self, is_train):
        super(MNIST, self).__init__()
        self.is_train = is_train

        if sys.platform == "win32":
            # 在windows平台下，使用dir命令获取路径下全部的包含.png的文件的路径
            # [:-1]去掉每行最后的回车符
            train_paths = [path[:-1] for path in os.popen('dir /s /b %s | findstr ".png"' % MNIST_TRAIN_PATH)._stream.readlines()]
            test_paths = [path[:-1] for path in os.popen('dir /s /b %s | findstr ".png"' % MNIST_TEST_PATH)._stream.readlines()]
        else:
            # 在Linux平台下，使用find命令获取路径下以.png结尾的文件的路径
            train_paths = [path[:-1] for path in os.popen('find %s | grep "\.png$"' % MNIST_TRAIN_PATH)._stream.readlines()]
            test_paths = [path[:-1] for path in os.popen('find %s | grep "\.png$"' % MNIST_TEST_PATH)._stream.readlines()]
        # 使用文件路径构造每个样本的记录
        self.train_records = [MNISTRecord(path) for path in train_paths]
        self.test_records = [MNISTRecord(path) for path in test_paths]

        # 打乱训练集和测试集的顺序
        random.seed(627)
        random.shuffle(self.train_records)
        random.shuffle(self.test_records)

    def x_transform(self, x):
        '''
        对读取到的图像进行预处理
        :param x:
        :return:
        '''
        # 将图像进行归一化，转为张量类型
        transform = torchvision.transforms.Compose([
                            ImageToTensor()])
        return transform(x)

    def __getitem__(self, item):
        if self.is_train:
            record = self.train_records[item]
        else:
            record = self.test_records[item]
        x = Image.open(record.path).convert('L')
        x = self.x_transform(x)

        return x, record.category


    def __len__(self):
        if self.is_train:
            return len(self.train_records)
        else:
            return len(self.test_records)

if __name__ == "__main__":
    dataset = MNIST(is_train=True)
    print(dataset[255])